iT邦幫忙

0

在 Django REST framework 裡按客戶做限流和每月額度:DRF throttling 不夠用的地方

  • 分享至 

  • xImage
  •  

利益揭露:我在做 UsageRail,一個幫 API 做 key、限流和用量計費的服務,最後一節會提到它。前面的程式都只用 Django 和 DRF,在 Django 6.1、DRF 3.18 上跑過。

API 開始收費以後,「這個 key 一分鐘能打幾次」和「這個客戶這個月還剩多少」會被放在同一張需求單上。兩件事長得很像,都是數請求、超過就回 429,但實際上差很多:

  • 限流(rate limit) 是保護伺服器的。窗口短(秒、分鐘),數字不準一點沒關係,Redis 重開歸零也無所謂。
  • 額度(quota) 是商業條件,例如免費方案每月 1,000 次。窗口長(一個帳期),數字要準,因為它最後會變成帳單上的一行。

這兩個東西該存的地方不一樣,該在什麼時候記一筆也不一樣。DRF 的 throttling 比較適合前者,硬拿來做後者會踩到幾個坑。下面從 DRF 內建的機制講起,再一步步補上按客戶的限流和額度。

一、DRF 內建 throttling 做了什麼

DRF 的 throttle 是在 APIView.initial() 裡跑的,順序是認證、權限、限流,全過了才進到你的 view。內建的 AnonRateThrottle、UserRateThrottle、ScopedRateThrottle 都繼承 SimpleRateThrottle,做法是在 Django cache 裡為每個識別碼存一串時間戳:讀出來、丟掉超過窗口的、數數看還剩幾格,沒超過就把這次的時間戳加上去,再寫回 cache。

拿來做「按客戶」的限流,有幾個地方不夠用:

  1. 速率是寫死在設定裡的。 DEFAULT_THROTTLE_RATES 是 {"user": "100/min"} 這種格式,一個 scope 一個數字。免費方案每分鐘 60 次、付費方案 600 次,沒辦法直接表達。SimpleRateThrottle.__init__() 會在建構時就讀設定、解析速率,那時候還不知道請求是誰發的。
  2. 讀和寫不是原子的。 從 cache 取出時間戳串、改完再寫回,兩個請求同時進來可能都讀到同一份,都覺得還有空位。DRF 官方文件也寫了內建實作有 race condition,高並發時會多放幾個請求過去。
  3. 預設 cache 是每個行程各一份。 沒設定 CACHES 時用的是 LocMemCache,gunicorn 開 4 個 worker,每個 worker 各數各的,限制等於變成 4 倍。
  4. 檢查時就記一筆。 時間戳是在 view 執行前寫進去的,所以回 400、500 的請求一樣佔名額。對限流來說這很合理(打爆伺服器的請求本來就該算),對額度和計費來說就不對了:客戶送了一個格式錯誤的 body,不應該扣他這個月的額度。
  5. 所有 throttle 都會跑完。 這點最容易忽略。DRF 的 check_throttles() 長這樣:
def check_throttles(self, request):
    throttle_durations = []
    for throttle in self.get_throttles():
        if not throttle.allow_request(request, self):
            throttle_durations.append(throttle.wait())
    ...

它不會在第一個失敗的 throttle 就停下,而是每個都問一遍,最後取最長的等待時間。如果你把「扣額度」寫成一個 throttle,跟限流 throttle 並排,那麼限流擋下來的請求,額度一樣會被扣掉。

所以接下來的做法是:限流留在 throttle 裡,換成按客戶、存在共用 cache;額度從 throttle 拿出來,放到資料庫,在限流通過之後才扣。

二、先把客戶認出來

限流和額度都需要知道「這是哪個客戶」。假設已經有這幾張表:

# billing/models.py
from django.db import models


class Plan(models.Model):
    name = models.CharField(max_length=50)
    # 留空代表不限
    requests_per_minute = models.PositiveIntegerField(null=True, blank=True)
    monthly_quota = models.PositiveIntegerField(null=True, blank=True)


class Customer(models.Model):
    name = models.CharField(max_length=100)
    plan = models.ForeignKey(Plan, on_delete=models.PROTECT)


class ApiKey(models.Model):
    customer = models.ForeignKey(Customer, on_delete=models.CASCADE, related_name="api_keys")
    key_id = models.CharField(max_length=16, unique=True)
    secret_hash = models.CharField(max_length=64)
    revoked_at = models.DateTimeField(null=True, blank=True)

認證類別把 key 對應到 ApiKey,放進 request.auth:

# billing/authentication.py
import hashlib
import hmac

from rest_framework import authentication, exceptions, permissions

from .models import ApiKey


class ApiKeyAuthentication(authentication.BaseAuthentication):
    def authenticate(self, request):
        header = request.headers.get("Authorization", "")
        if not header.startswith("Bearer "):
            return None
        # 假設 key 長這樣:{key_id}.{secret}
        key_id, _, secret = header.removeprefix("Bearer ").partition(".")
        api_key = (
            ApiKey.objects.select_related("customer__plan")
            .filter(key_id=key_id, revoked_at__isnull=True)
            .first()
        )
        secret_hash = hashlib.sha256(secret.encode()).hexdigest()
        if api_key is None or not hmac.compare_digest(api_key.secret_hash, secret_hash):
            raise exceptions.AuthenticationFailed("invalid_api_key")
        # 呼叫方是「客戶」,不是 Django 的 User,所以 user 留 None,客戶資料放在 request.auth
        return (None, api_key)

    def authenticate_header(self, request):
        # 有這行,沒帶 key 時 DRF 回 401 而不是 403
        return "Bearer"


class HasApiKey(permissions.BasePermission):
    def has_permission(self, request, view):
        return request.auth is not None

select_related("customer__plan") 不是順手加的。後面的 throttle 和額度都要讀 request.auth.customer.plan,沒有它,每個請求會多兩次查詢。

因為 user 是 None,內建的 IsAuthenticated 會一律拒絕,所以另外寫了 HasApiKey。

三、按客戶限流

做法 A:沿用 SimpleRateThrottle,只換速率來源

改動最少的方式,是讓 throttle 等拿到 request 之後才決定速率。DRF 自己的 ScopedRateThrottle 就是這樣做的:覆寫 __init__() 不讀設定,在 allow_request() 裡才設定 rate。

# billing/throttling.py
from rest_framework.throttling import SimpleRateThrottle


class PlanRateThrottle(SimpleRateThrottle):
    scope = "plan"

    def __init__(self):
        # 不在建構時讀 DEFAULT_THROTTLE_RATES,等拿到 request 才知道是哪個方案
        pass

    def allow_request(self, request, view):
        limit = request.auth.customer.plan.requests_per_minute if request.auth else None
        self.rate = f"{limit}/min" if limit else None
        self.num_requests, self.duration = self.parse_rate(self.rate)
        return super().allow_request(request, view)

    def get_cache_key(self, request, view):
        return self.cache_format % {"scope": self.scope, "ident": request.auth.customer_id}

識別碼用的是 customer_id,不是 key。一個客戶發了三把 key,三把加起來才是他的速率。

throttle 物件是每個請求新建一個的(get_throttles() 每次都 throttle()),所以在 allow_request() 裡改 self.rate 不會影響別的請求。

這個做法的限制就是第一節講的第 2、3 點:時間戳串的讀寫不是原子的,而且每個客戶在 cache 裡存一串長度等於上限的 list,上限 600 就是 600 個浮點數。

做法 B:固定窗口 + Redis INCR

如果要原子性,最簡單的是固定窗口計數:每分鐘一個 key,INCR 一下看有沒有超過。

# billing/throttling.py
import time

from django.core.cache import cache
from rest_framework.throttling import BaseThrottle


class CustomerRateThrottle(BaseThrottle):
    window_seconds = 60

    def allow_request(self, request, view):
        if request.auth is None:
            return True
        limit = request.auth.customer.plan.requests_per_minute
        if limit is None:
            return True
        now = time.time()
        window_start = int(now // self.window_seconds) * self.window_seconds
        cache_key = f"ratelimit:{request.auth.customer_id}:{window_start}"
        # add 只在 key 不存在時寫入;接著 incr 在 Redis 上是原子操作
        cache.add(cache_key, 0, timeout=self.window_seconds * 2)
        try:
            count = cache.incr(cache_key)
        except ValueError:
            # key 剛好在 add 和 incr 之間過期
            cache.set(cache_key, 1, timeout=self.window_seconds * 2)
            count = 1
        self.retry_after = window_start + self.window_seconds - now
        return count <= limit

    def wait(self):
        return self.retry_after

設定裡把 cache 換成 Redis(Django 4.0 起內建 RedisCache,不用另外裝 django-redis):

# settings.py
CACHES = {
    "default": {
        "BACKEND": "django.core.cache.backends.redis.RedisCache",
        "LOCATION": "redis://127.0.0.1:6379",
    }
}

REST_FRAMEWORK = {
    "DEFAULT_AUTHENTICATION_CLASSES": ["billing.authentication.ApiKeyAuthentication"],
    "DEFAULT_PERMISSION_CLASSES": ["billing.authentication.HasApiKey"],
    "DEFAULT_THROTTLE_CLASSES": ["billing.throttling.CustomerRateThrottle"],
    "EXCEPTION_HANDLER": "billing.handlers.api_exception_handler",
}

固定窗口的代價是邊界:上限 60 次的客戶,可以在 12:00:59 打 60 次、12:01:00 再打 60 次,兩秒內 120 次。如果你在意這個,可以改成 token bucket,用一段 Lua script 在 Redis 裡一次算完「補多少 token、扣一個、剩多少」;或者再加一個每秒的上限把突發壓住。只是要防止某個客戶的程式寫壞了狂打的話,固定窗口通常就夠了。

另外,這裡被擋下的請求也會讓計數加一。窗口一到就歸零,所以不會累積;但這代表一個持續狂打的客戶,在整個窗口內都會被擋著,這通常正是你要的。

四、每月額度要存在哪裡

額度不要放在 cache。理由有三個:

  • Redis 設了 maxmemory-policy allkeys-lru 之類的淘汰策略時,冷門客戶的計數可能被清掉,等於免費送他一個月。
  • Redis 重開或換機器,沒開持久化的話全部歸零。
  • 這個數字是帳單資料。客戶來問「為什麼我這個月被擋了」,你要能從資料庫查出來。

所以額度用一張表記,一個客戶一個帳期一列:

# billing/models.py
class UsageCounter(models.Model):
    customer = models.ForeignKey(Customer, on_delete=models.CASCADE)
    period_start = models.DateField()
    used = models.PositiveBigIntegerField(default=0)

    class Meta:
        constraints = [
            models.UniqueConstraint(fields=["customer", "period_start"], name="one_counter_per_period"),
        ]

扣額度的關鍵是把「還有沒有剩」和「加一」放在同一條 UPDATE:

# billing/quota.py
import datetime as dt

from django.db.models import F
from django.utils import timezone
from rest_framework.exceptions import Throttled

from .models import UsageCounter


class QuotaExceeded(Throttled):
    default_detail = "Monthly quota exceeded."
    default_code = "quota_exceeded"


def current_period():
    # 用 settings.TIME_ZONE 的日曆月,例如 Asia/Taipei
    start = timezone.localdate().replace(day=1)
    next_start = (start + dt.timedelta(days=32)).replace(day=1)
    return start, next_start


def reserve_unit(customer):
    start, next_start = current_period()
    counter, _ = UsageCounter.objects.get_or_create(customer=customer, period_start=start)
    rows = UsageCounter.objects.filter(pk=counter.pk)
    quota = customer.plan.monthly_quota
    if quota is not None:
        rows = rows.filter(used__lt=quota)
    # 條件和加一在同一條 UPDATE 裡,兩個請求同時進來也不會一起擠過最後一格
    if rows.update(used=F("used") + 1) == 0:
        reset_at = timezone.make_aware(dt.datetime.combine(next_start, dt.time.min))
        raise QuotaExceeded(wait=(reset_at - timezone.now()).total_seconds())
    return counter.pk


def release_unit(counter_pk):
    UsageCounter.objects.filter(pk=counter_pk, used__gt=0).update(used=F("used") - 1)

產生的 SQL 大概是 UPDATE ... SET used = used + 1 WHERE id = %s AND used < %s,回傳受影響的列數。0 列就是額度用完了。不用 select_for_update() 先鎖再改,也不用在 Python 裡比大小。

帳期的時區要先想清楚。timezone.localdate() 用的是 settings.TIME_ZONE,設成 Asia/Taipei 的話,每月 1 號台灣時間午夜重置。如果客戶是按 Stripe 訂閱週期付費,帳期應該跟著訂閱走(例如每月 17 號),這時 period_start 要從訂閱資料算,不能用日曆月。

這個設計的弱點是熱點列。一個每秒幾百個請求的大客戶,所有請求都在 UPDATE 同一列,PostgreSQL 的列鎖會讓它們排隊。量大到這個程度時,常見的做法是批次預扣:一次從資料庫扣 100 個單位,放在 Redis 或行程記憶體裡慢慢用,用完再扣下一批,結束時把沒用到的還回去。代價是要處理「行程掛掉時手上那批怎麼辦」。量沒到那裡之前,一條條件 UPDATE 就夠了。

五、計數什麼時候記

這是整件事最需要想清楚的地方。大致有三種時機:

時機 優點 問題
檢查時就記(DRF throttle 的做法) 簡單,超過就擋 400、500 也被扣
回應之後才記 只算成功的 並發時會超額:10 個請求同時檢查都看到「還剩 1」,全部放行
先預扣,結束時結算 不會超額,失敗的會還回去 要多寫一段結算

限流用第一種沒問題。額度和計費用第三種:進 view 之前先扣一格,回應是 4xx、5xx 就還回去。

在 DRF 裡,initial() 和 finalize_response() 剛好是這兩個掛點:

# billing/quota.py
class QuotaMixin:
    def initial(self, request, *args, **kwargs):
        # 先跑完認證、權限、限流,全部通過才扣額度
        super().initial(request, *args, **kwargs)
        self.quota_counter = reserve_unit(request.auth.customer)

    def finalize_response(self, request, response, *args, **kwargs):
        response = super().finalize_response(request, response, *args, **kwargs)
        counter = getattr(self, "quota_counter", None)
        if counter is not None and response.status_code >= 400:
            release_unit(counter)
        return response

用法:

from rest_framework.response import Response
from rest_framework.views import APIView

from billing.quota import QuotaMixin


class Generate(QuotaMixin, APIView):
    def post(self, request):
        ...
        return Response({"ok": True})

幾個細節:

  • super().initial() 跑完才扣,所以被限流擋下的請求不會扣到額度,第一節第 5 點的問題就沒了。
  • 扣額度失敗時 reserve_unit() 直接丟 QuotaExceeded,self.quota_counter 沒被設定,finalize_response() 不會誤還。
  • 結算用的是當初那一列的 pk,所以 1 月 31 日 23:59:59 進來、2 月 1 日才結束的請求,會還到 1 月那一列,不會還錯月份。
  • finalize_response() 只在 DRF 有處理到的例外路徑上會被呼叫。view 裡丟出 DRF 不認得的例外(例如一個普通的 KeyError)時,handle_exception() 會把它往外丟,finalize_response() 不會跑到,這一格就不會還。如果你很在意 500 不扣額度,可以覆寫 dispatch(),用 try/except 包起來自己還。

六、429 和 Retry-After

DRF 對 Throttled 例外有內建處理:狀態碼 429,wait 不是 None 就會加上 Retry-After 標頭(無條件進位成整數秒)。限流 throttle 的 wait() 回傳值、QuotaExceeded(wait=...) 傳進去的秒數,都會走這條路。

預設的 body 是 {"detail": "Request was throttled. Expected available in 54 seconds."},給人看可以,給程式判斷不太方便。用一個例外處理器把錯誤代碼放出來:

# billing/handlers.py
from rest_framework.exceptions import Throttled
from rest_framework.views import exception_handler


def api_exception_handler(exc, context):
    response = exception_handler(exc, context)
    if response is not None and isinstance(exc, Throttled):
        response.data = {"error": exc.get_codes(), "retry_after": exc.wait}
    return response

get_codes() 會回傳例外的 default_code,限流是 throttled,額度是 quota_exceeded。

這個區分很重要。限流的 Retry-After 是幾秒,額度的 Retry-After 是「到下個帳期還有幾秒」,月中被擋的話是一百多萬秒。用戶端如果寫了「收到 429 就照 Retry-After 睡一下再重試」,遇到後者會直接卡住好幾天。body 裡有 quota_exceeded,用戶端才知道這不是等一下就好,要升級方案或等下個月。

也有人額度用完回 402 或 403。402 比較適合預付點數用完(要付錢),403 容易跟權限錯誤混在一起。我自己傾向額度用完也回 429,用 body 的錯誤代碼區分。

七、跑一次看看

用上面的程式,建一個方案:每分鐘 5 次、每月 3 次。連打 7 次 POST /v1/generate,第一次故意送壞掉的 body:

#  狀態  Retry-After  body                                                  已用額度
0  400   -            {"error": "bad_input"}                                0
1  200   -            {"ok": true}                                          1
2  200   -            {"ok": true}                                          2
3  200   -            {"ok": true}                                          3
4  429   1386054      {"error": "quota_exceeded", "retry_after": 1386054}   3
5  429   54           {"error": "throttled", "retry_after": 54}             3
6  429   54           {"error": "throttled", "retry_after": 54}             3

可以看到:

  • 第 0 次回 400,額度被扣了又還回去,所以還是 0。
  • 第 4 次額度用完,Retry-After 是到下個月 1 號台灣時間午夜的秒數。
  • 第 5 次開始被限流擋下。0 到 4 這 5 次都佔了限流的名額,包括 400 和額度用完的那次,因為限流數的是「嘗試」,不是「成功」。
  • 被限流擋下的第 5、6 次,額度沒有被動到。

八、用 UsageRail 的 Django 中介軟體

上面的 key、限流、額度和用量計數都得自己維護,做法 B 還要一台 Redis,月底送帳和對帳也得另外寫。UsageRail 把這些拆成兩半:Django 中介軟體負責每個請求的檢查和計數,在 view 之前檢查 key、限流和額度,回應關閉之後才記用量,2xx、3xx 算計費、4xx 和 5xx 不算;送進 Stripe 和事後對帳則由 UsageRail 的服務端做。限流是每個客戶一個 token bucket,存在各行程的記憶體裡;額度則是由服務端把客戶剩下的額度切一小份給每台伺服器,伺服器在自己那份裡扣,大部分請求不用連網路,也就沒有第四節的熱點列問題。用量每 60 秒送一次,帳期結束後記到你自己的 Stripe 帳號開出的發票上。

# settings.py
import os

MIDDLEWARE = [
    # ...
    "django.contrib.auth.middleware.AuthenticationMiddleware",
    "usagerail.django.UsageRailMiddleware",
]

USAGERAIL = {
    "app_key": os.environ["USAGERAIL_APP_KEY"],
    "auth": "api-key",
}

跟 DRF 一起用時有兩點要注意:它在 DRF 之前就把請求擋下,所以 DRF 的例外處理器看不到這些 429,要改回應格式得用它的 on_reject;另外它讀的是 Authorization: Bearer 或 x-api-key,不讀 DRF TokenAuthentication 那種 Authorization: Token。設定方式、router 產生的 URL 怎麼記錄、行程被強制結束時還沒送出的用量會怎樣,都寫在 Django API 的計量與限流設定裡。


圖片
  熱門推薦
圖片
{{ item.channelVendor }} | {{ item.webinarstarted }} |
{{ formatDate(item.duration) }}
直播中

尚未有邦友留言

立即登入留言